/*
 * @Author: 李石
 * @Date: 2024-08-13 14:56:08
 * @LastEditors: lishi
 * @LastEditTime: 2024-08-13 17:46:29
 * @Description: 
 * Copyright (c) 2024 by ${lishi0105@163.com}, All Rights Reserved. 
 */
#include "rtsp_comm.h"
#include "nalu_parse.h"

#define NAL_REF_IDC_PRIORITY_HIGHEST    3
#define NAL_REF_IDC_PRIORITY_HIGH       2
#define NAL_REF_IDC_PRIORITY_LOW        1
#define NAL_REF_IDC_PRIORITY_DISPOSABLE 0

//Table 7-1 NAL unit type codes
#define NAL_UNIT_TYPE_UNSPECIFIED                    0    // Unspecified
#define NAL_UNIT_TYPE_CODED_SLICE_NON_IDR            1    // Coded slice of a non-IDR picture
#define NAL_UNIT_TYPE_CODED_SLICE_DATA_PARTITION_A   2    // Coded slice data partition A
#define NAL_UNIT_TYPE_CODED_SLICE_DATA_PARTITION_B   3    // Coded slice data partition B
#define NAL_UNIT_TYPE_CODED_SLICE_DATA_PARTITION_C   4    // Coded slice data partition C
#define NAL_UNIT_TYPE_CODED_SLICE_IDR                5    // Coded slice of an IDR picture
#define NAL_UNIT_TYPE_SEI                            6    // Supplemental enhancement information (SEI)
#define NAL_UNIT_TYPE_SPS                            7    // Sequence parameter set
#define NAL_UNIT_TYPE_PPS                            8    // Picture parameter set
#define NAL_UNIT_TYPE_AUD                            9    // Access unit delimiter
#define NAL_UNIT_TYPE_END_OF_SEQUENCE               10    // End of sequence
#define NAL_UNIT_TYPE_END_OF_STREAM                 11    // End of stream
#define NAL_UNIT_TYPE_FILLER                        12    // Filler data
#define NAL_UNIT_TYPE_SPS_EXT                       13    // Sequence parameter set extension
                                             // 14..18    // Reserved
#define NAL_UNIT_TYPE_CODED_SLICE_AUX               19    // Coded slice of an auxiliary coded picture without partitioning
                                             // 20..23    // Reserved
                                             // 24..31    // Unspecified
enum NalUnitType
{
    NAL_UNIT_CODED_SLICE_TRAIL_N = 0, // 0
    NAL_UNIT_CODED_SLICE_TRAIL_R,     // 1

    NAL_UNIT_CODED_SLICE_TSA_N,       // 2
    NAL_UNIT_CODED_SLICE_TSA_R,       // 3

    NAL_UNIT_CODED_SLICE_STSA_N,      // 4
    NAL_UNIT_CODED_SLICE_STSA_R,      // 5

    NAL_UNIT_CODED_SLICE_RADL_N,      // 6
    NAL_UNIT_CODED_SLICE_RADL_R,      // 7

    NAL_UNIT_CODED_SLICE_RASL_N,      // 8
    NAL_UNIT_CODED_SLICE_RASL_R,      // 9

    NAL_UNIT_RESERVED_VCL_N10,
    NAL_UNIT_RESERVED_VCL_R11,
    NAL_UNIT_RESERVED_VCL_N12,
    NAL_UNIT_RESERVED_VCL_R13,
    NAL_UNIT_RESERVED_VCL_N14,
    NAL_UNIT_RESERVED_VCL_R15,

    NAL_UNIT_CODED_SLICE_BLA_W_LP,    // 16
    NAL_UNIT_CODED_SLICE_BLA_W_RADL,  // 17
    NAL_UNIT_CODED_SLICE_BLA_N_LP,    // 18
    NAL_UNIT_CODED_SLICE_IDR_W_RADL,  // 19
    NAL_UNIT_CODED_SLICE_IDR_N_LP,    // 20
    NAL_UNIT_CODED_SLICE_CRA,         // 21
    NAL_UNIT_RESERVED_IRAP_VCL22,
    NAL_UNIT_RESERVED_IRAP_VCL23,

    NAL_UNIT_RESERVED_VCL24,
    NAL_UNIT_RESERVED_VCL25,
    NAL_UNIT_RESERVED_VCL26,
    NAL_UNIT_RESERVED_VCL27,
    NAL_UNIT_RESERVED_VCL28,
    NAL_UNIT_RESERVED_VCL29,
    NAL_UNIT_RESERVED_VCL30,
    NAL_UNIT_RESERVED_VCL31,

    // non-VCL
    NAL_UNIT_VPS,                     // 32
    NAL_UNIT_SPS,                     // 33
    NAL_UNIT_PPS,                     // 34
    NAL_UNIT_AUD,                     // 35
    NAL_UNIT_EOS,                     // 36
    NAL_UNIT_EOB,                     // 37
    NAL_UNIT_FILLER_DATA,             // 38
    NAL_UNIT_PREFIX_SEI,              // 39
    NAL_UNIT_SUFFIX_SEI,              // 40

    NAL_UNIT_RESERVED_NVCL41,
    NAL_UNIT_RESERVED_NVCL42,
    NAL_UNIT_RESERVED_NVCL43,
    NAL_UNIT_RESERVED_NVCL44,
    NAL_UNIT_RESERVED_NVCL45,
    NAL_UNIT_RESERVED_NVCL46,
    NAL_UNIT_RESERVED_NVCL47,
    NAL_UNIT_UNSPECIFIED_48,
    NAL_UNIT_UNSPECIFIED_49,
    NAL_UNIT_UNSPECIFIED_50,
    NAL_UNIT_UNSPECIFIED_51,
    NAL_UNIT_UNSPECIFIED_52,
    NAL_UNIT_UNSPECIFIED_53,
    NAL_UNIT_UNSPECIFIED_54,
    NAL_UNIT_UNSPECIFIED_55,
    NAL_UNIT_UNSPECIFIED_56,
    NAL_UNIT_UNSPECIFIED_57,
    NAL_UNIT_UNSPECIFIED_58,
    NAL_UNIT_UNSPECIFIED_59,
    NAL_UNIT_UNSPECIFIED_60,
    NAL_UNIT_UNSPECIFIED_61,
    NAL_UNIT_UNSPECIFIED_62,
    NAL_UNIT_UNSPECIFIED_63,
    NAL_UNIT_INVALID,
};

#define H264_TYPE(v) ((uint8_t)(v) & 0x1F)
#define H265_TYPE(v) (((uint8_t)(v) >> 1) & 0x3f)

static int32_t _startLen(const uint8_t *in)
{
    int32_t start_ = 0;
    if(in[0] == 0x0){
        if(in[2] == 0x01) start_ = 3;
        else if(in[2] == 0x00 && in[3] == 0x01) start_ = 4;
        else start_ = -1;
    }
    return start_;
}

char LibRtsp_keyframe(int type, const uint8_t *in, size_t size, int32_t *startLen_)
{
    *startLen_ = 0;
    if(in == NULL || size < 5)return NaluFalse;
    int32_t start_ = _startLen(in);
    if(start_ == -1)return NaluFalse;
    *startLen_ = start_;
    if(type == LIBRTSP_VIDEO_H264){
        int8_t nalu_type = H264_TYPE(in[start_]);
        return (nalu_type == 
            NAL_UNIT_TYPE_CODED_SLICE_IDR) && LibRtsp_decodeAble(type, in, size, startLen_);
    }else if(type == LIBRTSP_VIDEO_H265){
        int8_t nalu_type = H265_TYPE(in[start_]);
        return ((nalu_type >= NAL_UNIT_CODED_SLICE_BLA_W_LP && 
                 nalu_type <= NAL_UNIT_RESERVED_IRAP_VCL23) && 
                 LibRtsp_decodeAble(type, in, size, startLen_));
    }
    return NaluFalse;
}

char LibRtsp_configFrame(int type, const uint8_t *in, size_t size, int32_t *startLen_)
{
    *startLen_ = 0;
    if(in == NULL || size < 5)return NaluFalse;
    int32_t start_ = _startLen(in);
    if(start_ == -1)return NaluFalse;
    *startLen_ = start_;
    if(type == LIBRTSP_VIDEO_H264){
        int8_t nalu_type = H264_TYPE(in[start_]);
        if(nalu_type == NAL_UNIT_TYPE_SPS || 
           nalu_type == NAL_UNIT_TYPE_PPS) return NaluTrue;
    }else if(type == LIBRTSP_VIDEO_H265){
        int8_t nalu_type = H265_TYPE(in[start_]);
        if(nalu_type == NAL_UNIT_VPS || 
           nalu_type == NAL_UNIT_SPS || 
           nalu_type == NAL_UNIT_PPS) return NaluTrue;
    }
    return NaluFalse;
}

char LibRtsp_dropAble(int type, const uint8_t *in, size_t size, int32_t *startLen_)
{
    *startLen_ = 0;
    if(in == NULL || size < 5)return NaluFalse;
    int32_t start_ = _startLen(in);
    if(start_ == -1)return NaluFalse;
    *startLen_ = start_;
    if(type == LIBRTSP_VIDEO_H264){
        int8_t nalu_type = H264_TYPE(in[start_]);
        if(nalu_type == NAL_UNIT_TYPE_SEI || 
           nalu_type == NAL_UNIT_TYPE_AUD) return NaluTrue;
    }else if(type == LIBRTSP_VIDEO_H265){
        int8_t nalu_type = H265_TYPE(in[start_]);
        if(nalu_type == NAL_UNIT_PREFIX_SEI || 
           nalu_type == NAL_UNIT_SUFFIX_SEI || 
           nalu_type == NAL_UNIT_AUD) return NaluTrue;
    }
    return NaluFalse;
}

char LibRtsp_decodeAble(int type, const uint8_t *in, size_t size, int32_t *startLen_)
{
    *startLen_ = 0;
    if(in == NULL || size < 5)return NaluFalse;
    int32_t start_ = _startLen(in);
    if(start_ == -1)return NaluFalse;
    *startLen_ = start_;
    if(type == LIBRTSP_VIDEO_H264){
        int8_t nalu_type = H264_TYPE(in[start_]);
        return (nalu_type >= NAL_UNIT_TYPE_CODED_SLICE_NON_IDR && 
                nalu_type <= NAL_UNIT_TYPE_CODED_SLICE_IDR && 
                (in[start_+1] & 0x80));
    }else if(type == LIBRTSP_VIDEO_H265){
        int8_t nalu_type = H265_TYPE(in[start_]);
        return (nalu_type >= NAL_UNIT_CODED_SLICE_TRAIL_N && 
                nalu_type <= NAL_UNIT_RESERVED_IRAP_VCL23 && 
                (in[start_+2] & 0x80));
    }
    return NaluFalse;
}

//分离sps/pps/vps配置帧和I帧在一起的组合帧
char LibRtsp_separateFrame(int type, const void* data, size_t size, LibRtspNalus *nalus)
{
    size_t i = 0, last_nalu_pos = 0;
    const uint8_t *in = (const uint8_t *)data;
    int32_t start_len = 0;
    uint32_t nalu_size = 0;
    if(nalus == NULL){
        rtsp_err("nalus is NULL\n");
        return NaluFalse;
    }
    if(data == NULL || size < 5){
        rtsp_err("data is NULL or size:(%lu)\n", size);
        return NaluFalse;
    }
    if(type != LIBRTSP_VIDEO_H264 && type != LIBRTSP_VIDEO_H265){
        rtsp_err("type:(%d) error\n", type);
        return NaluFalse;
    }
    if( !LibRtsp_configFrame(type, in, size, &start_len) || 
        LibRtsp_dropAble(type, in, size, &start_len)){
        // 如果起始NALU单元为非配置帧或可丢弃帧，则认为只包含一个NALU单元
        nalus->count = 1;
        nalus->nalus[0].type = type;
        nalus->nalus[0].num = 0;
        nalus->nalus[0].len = size;
        nalus->nalus[0].offset = 0;
        nalus->nalus[0].startcodeLen = start_len;
        return NaluTrue;
    }
    for (i = 0; i < size; i++) {
        if (i < size - 4 &&
            in[i] == 0x00 && in[i + 1] == 0x00 &&
            ((in[i + 2] == 0x01) || (in[i + 2] == 0x00 && in[i + 3] == 0x01))) {
            int startcodeLen = (in[i + 2] == 0x01) ? 2 : 3;
            if (nalu_size > 0) {
                nalus->nalus[nalu_size - 1].len = i - last_nalu_pos;
            }

            // Record the current NALU header position
            last_nalu_pos = i;
            nalus->nalus[nalu_size].offset = last_nalu_pos;
            nalus->nalus[nalu_size].startcodeLen = startcodeLen + 1;
            nalus->nalus[nalu_size].type = type;
            nalus->nalus[nalu_size].num = nalu_size;
            i += startcodeLen;
            nalu_size ++;
            if(nalu_size >= FRAME_MAX_NALU_COUNT){
                break;
            }
            if( !LibRtsp_configFrame(type, in + last_nalu_pos, size - last_nalu_pos, &start_len) &&  
                LibRtsp_decodeAble(type, in + last_nalu_pos, size - last_nalu_pos, &start_len)){
                // 如果是可解码的非配置帧，则认为nalu单元已全部解析完毕
                break;
            }
        }
    }
    if (nalu_size > 0) {
        nalus->nalus[nalu_size - 1].len = size - last_nalu_pos;
    }
    nalus->count = nalu_size;
    return NaluTrue;
}